﻿/*
 * OscarLib
 * http://shaim.net/trac/oscarlib/
 * Copyright ©2005-2008, Chris Sammis
 * Licensed under the Lesser GNU Public License (LGPL)
 * http://www.opensource.org/osi3.0/licenses/lgpl-license.php
 * 
 */

using System;
using System.Collections.Generic;
using System.Text;
using System.Net.Sockets;
using System.Net;
using System.IO;

namespace csammisrun.OscarLib.Utility
{
    /// <summary>
    /// Encapsulates connecting to a server through a proxy
    /// </summary>
    class ProxySocketConnector
    {
        private readonly string destinationServer;
        private readonly int destinationPort;
        private readonly Session parent;

        private Socket socket;

        /// <summary>
        /// Raised when the connection is complete and the socket is ready for use
        /// </summary>
        public event EventHandler ConnectionComplete;
        /// <summary>
        /// Raised when the socket could not be connected as requested
        /// </summary>
        public event EventHandler ConnectionFailed;

        /// <summary>
        /// Creates a new ProxySocketConnector
        /// </summary>
        /// <param name="parent">The <see cref="Session"/> containing proxy settings to use</param>
        /// <param name="destinationServer">The destination server</param>
        /// <param name="destinationPort">The destination port</param>
        public ProxySocketConnector(Session parent, string destinationServer, int destinationPort)
        {
            this.parent = parent;
            this.destinationServer = destinationServer;
            this.destinationPort = destinationPort;
        }

        /// <summary>
        /// Begins the connection process asynchronously
        /// </summary>
        public void BeginConnect()
        {
            string serverToUse = GetServerForInitialConnection();

            IPAddress ipAddress = null;
            if (IPAddress.TryParse(serverToUse, out ipAddress))
            {
                ConnectToIPAddress(ipAddress);
            }
            else
            {
                try
                {
                    Dns.BeginGetHostEntry(serverToUse, EndGetHostEntry, null);
                }
                catch (SocketException sockex)
                {
                    OnConnectionFailed("Exception in GetHostEntry: {0}", sockex);
                }
            }
        }

        /// <summary>
        /// Gets the <see cref="Socket"/> generated by the connector
        /// </summary>
        public Socket Socket
        {
            get { return socket; }
        }

        #region Protected methods
        /// <summary>
        /// Raises the <see cref="ConnectionFailed"/> event
        /// </summary>
        protected void OnConnectionFailed(string message, params object[] parameters)
        {
            Logging.WriteString(message, parameters);
            if (ConnectionFailed != null)
            {
                ConnectionFailed(this, new EventArgs());
            }
        }

        /// <summary>
        /// Raises the <see cref="ConnectionComplete"/> event
        /// </summary>
        protected void OnConnectionComplete()
        {
            if (ConnectionComplete != null)
            {
                ConnectionComplete(this, new EventArgs());
            }
        }

        #endregion Protected methods

        private void EndGetHostEntry(IAsyncResult res)
        {
            IPHostEntry hosts = null;
            try
            {
                hosts = Dns.EndGetHostEntry(res);
            }
            catch (SocketException)
            {
                // This is handled by the fallback DNS lookup below
            }

            if (hosts == null || hosts.AddressList == null || hosts.AddressList.Length == 0)
            {
                Logging.WriteString("Using fallback DNS lookup method to connect to server");

                string errorMessage = null;

                // Fallback to a different synchronous DNS lookup method
                try
                {
                    hosts = new IPHostEntry();
                    hosts.AddressList = Dns.GetHostAddresses(GetServerForInitialConnection());
                    if (hosts.AddressList == null || hosts.AddressList.Length == 0)
                    {
                        throw new InvalidDataException();
                    }
                }
                catch (SocketException sockex)
                {
                    errorMessage = String.Format("Cannot resolve server: {0}", sockex);
                }
                catch (InvalidDataException)
                {
                    errorMessage = "Cannot resolve server: Dns.GetHostAddresses returned no entries";
                }

                if (!String.IsNullOrEmpty(errorMessage))
                {
                    OnConnectionFailed(errorMessage);
                    return;
                }
            }

            ConnectToIPAddress(hosts.AddressList[0]);
        }

        /// <summary>
        /// Connects directly to an IP address
        /// </summary>
        /// <param name="address">The IP address to which to connect</param>
        private void ConnectToIPAddress(IPAddress address)
        {
            IPEndPoint ipep = new IPEndPoint(address, GetPortForInitialConnection());
            socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
            try
            {
                socket.BeginConnect(ipep, EndConnectToIPAddress, null);
            }
            catch (SocketException sockex)
            {
                OnConnectionFailed("Cannot connect to server: {0}", sockex);
            }
        }

        /// <summary>
        /// Ends the host connection phase and negotiates the proxied connection
        /// </summary>
        private void EndConnectToIPAddress(IAsyncResult res)
        {
            try
            {
                socket.EndConnect(res);
            }
            catch (SocketException sockex)
            {
                OnConnectionFailed("Can't connect to server: {0}", sockex);
                return;
            }

            if (parent.ProxySetting == ProxyType.None)
            {
                OnConnectionComplete();
                return;
            }


            if (parent.ProxySetting == ProxyType.Socks5)
            {
                Socks5AuthType auth = GetSocks5AuthMethod();
                if (auth == Socks5AuthType.AuthMethodRejected)
                {
                    socket.Close();
                    OnConnectionFailed("SOCKS5 authentication mechanism unsupported");
                    return;
                }
                else if (auth != Socks5AuthType.None)
                {
                    if (SendSocks5Authentication() == false)
                    {
                        socket.Close();
                        OnConnectionFailed("SOCK5 authentication failed, bad username/password");
                        return;
                    }
                }

                if (SendSocks5ConnectRequest() == false)
                {
                    socket.Close();
                    OnConnectionFailed("SOCKS5 connect failed, proxy server failure");
                    return;
                }
            }

            // Done!
            OnConnectionComplete();
        }

        #region SOCK5 specific connection methods
        /// <summary>
        /// Queries the SOCKS5 proxy for its desired authentication mechanism
        /// </summary>
        private Socks5AuthType GetSocks5AuthMethod()
        {
            Socks5AuthType retval = Socks5AuthType.AuthMethodRejected;
            int reqlength = 4, sent = 0;
            byte[] request = new byte[reqlength];
            request[0] = 0x05;
            request[1] = 0x02;
            request[2] = 0x00;
            request[3] = 0x02;

            try
            {
                while (sent < reqlength)
                {
                    sent += socket.Send(request, sent, reqlength - sent, SocketFlags.None);
                }

                int replength = 2, received = 0;
                byte[] reply = new byte[replength];
                while (received < replength)
                {
                    received += socket.Receive(reply, received, replength - received, SocketFlags.None);
                }

                retval = (Socks5AuthType)reply[1];
            }
            catch (SocketException sockex)
            {
                Logging.WriteString("Caught exception in GetSocks5AuthMethod: {0}", sockex);
            }
            return retval;
        }

        /// <summary>
        /// Sends the authentication sequence to the proxy server
        /// </summary>
        private bool SendSocks5Authentication()
        {
            bool retval = false;
            int usernamelen = Encoding.ASCII.GetByteCount(parent.ProxyUsername);
            int passwordlen = Encoding.ASCII.GetByteCount(parent.ProxyPassword);
            int reqlength = 3 + usernamelen + passwordlen, sent = 0;
            byte[] request = new byte[reqlength];
            request[0] = 0x05;
            request[1] = (byte)usernamelen;
            Array.ConstrainedCopy(Encoding.ASCII.GetBytes(parent.ProxyUsername), 0, request, 2, usernamelen);
            request[2 + usernamelen] = (byte)passwordlen;
            Array.ConstrainedCopy(Encoding.ASCII.GetBytes(parent.ProxyPassword), 0, request, 3 + usernamelen, passwordlen);

            try
            {
                while (sent < reqlength)
                {
                    sent += socket.Send(request, sent, reqlength - sent, SocketFlags.None);
                }

                int replength = 2, received = 0;
                byte[] reply = new byte[replength];
                while (received < replength)
                {
                    received += socket.Receive(reply, received, replength - received, SocketFlags.None);
                }
                retval = (reply[1] == 0x00);
            }
            catch (SocketException sockex)
            {
                Logging.WriteString("Caught exception in SendSocks5Authentication: {0}", sockex);
            }
            return retval;
        }

        /// <summary>
        /// Send a CONNECT request to the proxy server
        /// </summary>
        private bool SendSocks5ConnectRequest()
        {
            bool retval = false;

            IPAddress ipaddr;
            if (!IPAddress.TryParse(destinationServer, out ipaddr))
            {
                ipaddr = null;
            }

            int sent = 0;
            int iplength = (ipaddr == null) ? Encoding.ASCII.GetByteCount(destinationServer) :
              (ipaddr.AddressFamily == AddressFamily.InterNetwork) ? 4 : 6;

            int reqlength = 6 + iplength;
            byte[] request = new byte[reqlength];
            request[0] = 0x05;
            request[1] = 0x01;
            request[2] = 0x00;
            if (ipaddr == null)
            {
                request[3] = 0x03;
                Array.ConstrainedCopy(Encoding.ASCII.GetBytes(destinationServer), 0, request, 4, reqlength - 6);

            }
            else if (ipaddr.AddressFamily == AddressFamily.InterNetwork)
            {
                request[3] = 0x01;
                Array.ConstrainedCopy(ipaddr.GetAddressBytes(), 0, request, 4, reqlength - 6);
            }
            else
            {
                request[3] = 0x04;
                Array.ConstrainedCopy(ipaddr.GetAddressBytes(), 0, request, 4, reqlength - 6);
            }

            request[reqlength - 2] = (byte)((destinationPort & 0xFF00) >> 8);
            request[reqlength - 1] = (byte)(destinationPort & 0xFF);
            try
            {
                while (sent < reqlength)
                {
                    sent += socket.Send(request, sent, reqlength - sent, SocketFlags.None);
                }

                byte[] reply = new byte[256]; // Variable length reply here
                socket.Receive(reply);
                retval = (reply[1] == 0x00);
            }
            catch (SocketException sockex)
            {
                Logging.WriteString("Caught exception in SendSocks5ConnectRequest: {0}", sockex);
            }
            return retval;
        }
        #endregion SOCKS5 specific connection methods

        /// <summary>
        /// Gets the server to connect to initially
        /// </summary>
        private string GetServerForInitialConnection()
        {
            string serverToUse = destinationServer;
            if (parent.ProxySetting != ProxyType.None)
            {
                serverToUse = parent.ProxyServer;
            }
            return serverToUse;
        }

        /// <summary>
        /// Gets the port to connect to initially
        /// </summary>
        private int GetPortForInitialConnection()
        {
            int portToUse = destinationPort;
            if (parent.ProxySetting != ProxyType.None)
            {
                portToUse = parent.ProxyPort;
            }
            return portToUse;
        }
    }
}
